# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: xguo

import asyncio
import datetime
import functools
import logging
import os
import signal
import time

from absl import app, flags
from coin.base.timestamp import get_timestamp
from tornado.ioloop import IOLoop
from datatool import to_next_mid_night

from coin.experimental.xguo.fruit.util.util import (
    guess_tick_info, calculate_precision)
from coin.experimental.xguo.fruit.util.record_writer import RecordWriter
from coin.exchange.base.rest.task_scheduler import Scheduler
from coin.exchange.base.rest_order_gateway_base import pretty_json
from coin.exchange.base.strategy_logger import StrategyLogProvider
from coin.experimental.xguo.fruit.base.strategy_config import BaseStrategyConfig, construct_executor
from coin.flow.simple_queue import SimpleQueue
from coin.flow.subscriber import from_queue
from coin.strategy.executor.util.cancel_all_executor import SimpleCancelAllExecutor
from coin.strategy.mm.base.bookmap import BookMap
from coin.strategy.mm.base.multi_feed_subsystem import MultiFeedSubsystem
from coin.strategy.mm.subscription import SubscriptionRequest
from coin.support.accounting.checker.account_checker import AccountChecker
from coin.base.book.util import (
    convert_book_to_dict,
    convert_trade_to_dict,
    human_readable_timestamp_record,
)
from coin.experimental.xguo.fruit.base.event_clock import (
    event_clock,)

from coin.strategy.mm.live_envs import (create_live_env_instance, get_live_env_cls)

from coin.exchange.bithumb.kr_rest.product_info import bithumb_tick_info
from coin.exchange.coinone_v2.kr_rest.product_info import coinone_tick_info
from coin.exchange.upbit_v1.kr_rest.product_info import upbit_tick_info


FLAGS = flags.FLAGS


def get_dryrun():
  dry_run = os.getenv('DRY_RUN')

  if dry_run is None:
    return False

  assert isinstance(dry_run, str), dry_run
  dry_run = dry_run.upper()
  if dry_run in ('', '0', 'NO', 'FALSE'):
    return False
  else:
    return dry_run


class BaseStrategy(object):
  def __init__(self, strategy_config: BaseStrategyConfig):
    self._logger = logging.getLogger(__name__)
    self._ioloop = IOLoop.current()
    self._queue = SimpleQueue('strategy')
    self._flow_sub = from_queue(self._queue)
    self._strategy_config = strategy_config
    self._feed_subsystem = None
    self._order_subsystems = None
    self._executors = None
    self._balance_printer = None
    self._balance_checker = None
    self._scheduler = Scheduler(ioloop=self._ioloop, logger=self._logger)
    self._name = None
    self._revision = None
    self._start_time = None
    self._book_map = BookMap(self.product_list)
    self._writer = RecordWriter(filename="smm_record_%s.gz" % int(time.time() * 1e9))

    core_config = self.strategy_config.strategy_core_config
    strategy_request = core_config['strategy_request']
    self._strategy_logger = StrategyLogProvider(strategy_request)
    self._dry_run = get_dryrun()
    self._executor_init_done = False
    self._guess_tick_info_done = False

  def is_order_subsystem_ready(self):
    if self._order_subsystems is None:
      return None

    is_ready = all(
        order_subsystem.ready_flag for order_subsystem in self._order_subsystems.values())

    for key, og in self._order_subsystems.items():
      self._logger.info("OG: %s, %s, %s\n", key, og, og.ready_flag)
    return is_ready

  def is_initialized(self):
    return self._guess_tick_info_done and self._executor_init_done

  @property
  def strategy_config(self):
    return self._strategy_config

  @property
  def order_gateway_list(self):
    if not self.is_order_subsystem_ready():
      return []

    if self._order_subsystems is None:
      return []

    og_list = [order_subsystem.order_gateway for order_subsystem in self._order_subsystems.values()]
    return og_list

  @property
  def executor_list(self):
    if not self.is_initialized():
      return []

    assert isinstance(self._executors, dict), self._executors
    return list(self._executors.values())

  @property
  def product_list(self):
    config_dict = self._strategy_config.feed_subsystem_config
    products = []
    if config_dict is None:
      self._logger.info('Feed system config not found!')
      return products

    for config in config_dict.values():
      products.extend(config.products)
    return products

  def get_executor(self, product):
    return self._executors[product]

  def prepare_feed_subsystem(self):
    config_dict = self._strategy_config.feed_subsystem_config
    if config_dict is None:
      self._logger.info('Feed system config not found!')
      return

    self._logger.info('Prepare feed system.')
    self._feed_subsystem = MultiFeedSubsystem(
        feed_checker_config=None,
        feed_checker_config_map=None,
    )

    for product in self.product_list:
      self._feed_subsystem.subscribe_book(product, functools.partial(self.common_on_book, product))
      self._feed_subsystem.subscribe_trade(product, functools.partial(self.on_trade, product))
    self._feed_subsystem.start(self._ioloop, self._queue, self._flow_sub)

  def prepare_order_system(self):
    config_list = self._strategy_config.order_subsystem_config
    if config_list is None:
      self._logger.info('Order system config not found!')
      return

    if self._dry_run:
      self._logger.info('Run in dryrun mode! OG will not be created!')
      return

    self._logger.info('Prepare order system.')
    self._order_subsystems = {}
    init_og_after = datetime.timedelta(seconds=5)
    for config in config_list.values():
      order_subsystem_cls = get_live_env_cls(config.subreq)
      order_subsystem = create_live_env_instance(
          order_subsystem_cls,
          currencies=config.currencies,
          trade_products=config.products,
          og_trade_key=config.og_config,
          og_params=config.og_params,
          init_og_after=init_og_after,
          ioloop=self._ioloop,
          queue=self._queue,
          flow_sub=self._flow_sub,
          og_log_callback=self._strategy_logger.write_og_log,
      )
      order_subsystem.ready_flag = False

      # noinspection PyUnusedLocal
      def set_ready_flag(subreq, order_gateway):
        order_subsystem.ready_flag = True

      order_subsystem.on_order_gateway_reset_callback.append(set_ready_flag)
      self._order_subsystems[config.subreq] = order_subsystem
      order_subsystem.start()

  async def prepare_executors(self):
    if self._strategy_config.order_subsystem_config is None:
      return

    if self._dry_run:
      self._logger.info('Run in dryrun mode, no live executors are created!')
      self._executor_init_done = True
      if self._guess_tick_info_done:
        self.run_before_start()
      return

    while not self.is_order_subsystem_ready():
      self._logger.warning('Order subsystem not ready yet! Wait for 1 seconds')
      await asyncio.sleep(1)

    config_list = self._strategy_config.order_subsystem_config
    if config_list is None:
      self._logger.info('Executor config not found!')
      return

    self._logger.info('Prepare executors')
    self._executors = {}
    logger = self._logger
    for key, config in self._strategy_config.executor_config.items():
      self._logger.info('Prepare %s executor', key)
      # Key is usually product
      subreq = SubscriptionRequest.from_str(config['subreq'])
      assert subreq
      proto_config = config['proto_config']
      og = self._order_subsystems[subreq].order_gateway
      self._executors[key] = construct_executor(subreq, proto_config, og, logger)
    for og in self.order_gateway_list:
      og.register_fill_callback(self.on_fill)
      self._logger.info('Register fill callback: %s', og)

    self._executor_init_done = True
    if self._guess_tick_info_done:
      self.run_before_start()

  def prepare_balance_checker(self):
    config = self._strategy_config.account_checker_config
    if config is None:
      self._logger.info('Balance checker config not found!')
      return

    self._logger.info('Prepare balance checker')
    self._balance_checker = AccountChecker(self._strategy_config.account_checker_config)

  def log_json_message(self, data, *, msg='', level='INFO', pretty=True):
    assert isinstance(data, (list, dict)), type(data)
    if isinstance(level, str):
      level = logging.getLevelName(level)
    self._logger.log(level, '%s:\n%s\n', msg, pretty_json(data))

  def on_book(self, product, book):
    book_dict = convert_book_to_dict(book)
    self.log_json_message(book_dict, msg='book')

  def common_on_book(self, product, book):
    self._book_map.set_book(product, book)
    event_clock.update(book.timestamp)

    if not self._guess_tick_info_done:
      if product.exchange == 'Bithumb':
        bithumb_tick_info[product.symbol] = guess_tick_info(book)
        self._logger.info("bithumb tick info: %s", bithumb_tick_info)
      elif product.exchange == 'Coinone':
        coinone_tick_info[product.symbol] = guess_tick_info(book)
        self._logger.info("coinone tick info: %s", coinone_tick_info)
      elif product.exchange == 'Upbit':
        upbit_tick_info[product.symbol] = guess_tick_info(book)
        self._logger.info("upbit tick info: %s", upbit_tick_info)
      else:
        pass

      num_bithumb = 0
      num_coinone = 0
      num_upbit = 0
      for prod in self.product_list:
        if prod.exchange == 'Bithumb':
          num_bithumb += 1
        elif prod.exchange == 'Coinone':
          num_coinone += 1
        elif prod.exchange == 'Upbit':
          num_upbit += 1
      a = num_bithumb == len(bithumb_tick_info)
      b = num_coinone == len(coinone_tick_info)
      c = num_upbit == len(upbit_tick_info)
      if a and b and c:
        self._guess_tick_info_done = True
        if self._executor_init_done:
          self.run_before_start()

    self.on_book(product, book)

  def on_trade(self, product, trade_event):
    # noinspection PyProtectedMember
    trade_dict = convert_trade_to_dict(product, trade_event)
    self.log_json_message(trade_dict, msg='trade')

  def on_fill(self, price_last_fill, qty_last_fill, order):
    # noinspection PyProtectedMember
    fill_dict = {
        'price_last_fill': price_last_fill,
        'qty_last_fill': qty_last_fill,
        'order': order._asdict()
    }
    self.log_json_message(fill_dict, msg='fill')

  def has_pending_orders(self):
    ogs = self.order_gateway_list
    return sum(len(og.get_working_order()) for og in ogs) > 0

  def cancel_all(self):
    exit_executor = SimpleCancelAllExecutor()
    ogs = self.order_gateway_list
    exit_executor.cancel_all_on_exit(ogs)

  def signal_handler(self, sig, frame):
    self._logger.warning('Caught signal: %s (%s)', sig, frame)
    self._ioloop.add_callback_from_signal(self.stop)

  def run(self):
    self.prepare_feed_subsystem()
    self.prepare_order_system()
    self.prepare_balance_checker()
    self._ioloop.add_callback(self.prepare_executors)

    signal.signal(signal.SIGTERM, self.signal_handler)
    signal.signal(signal.SIGINT, self.signal_handler)
    self._start_time = get_timestamp()
    data = {'type': 'strategy_started', 'data': {'msg': 'strategy started'}}
    data['data'].update(human_readable_timestamp_record(self._start_time))
    self.log_json_message(data, msg='STRATEGY_STARTED')
    next_mid_night = to_next_mid_night(self._start_time)
    self._ioloop.call_at(next_mid_night.timestamp(), self.stop)
    self._ioloop.start()

  def run_before_start(self):
    pass

  def stop(self):
    data = {'type': 'strategy_stopped', 'data': {'msg': 'strategy stopped'}}
    end_time = get_timestamp()
    data['data'].update(human_readable_timestamp_record(end_time))
    self.log_json_message(data, msg='STRATEGY_STOPPED')
    self._ioloop.add_callback(self._try_clean_stop)

  async def _try_clean_stop(self):
    # noinspection PyBroadException
    try:
      for _ in range(10):
        if self.has_pending_orders():
          self.cancel_all()
          await asyncio.sleep(3)
    except Exception:
      self._logger.info('Failed to clear pending orders!')
    self._ioloop.add_callback(self._ioloop.stop)


def main(_):
  if FLAGS.config is None:
    config_file = os.path.join(os.path.dirname(__file__), 'config.yaml')
  else:
    config_file = FLAGS.config
  config = BaseStrategyConfig.from_file(config_file)
  strategy = BaseStrategy(strategy_config=config)
  strategy.run()


if __name__ == '__main__':
  flags.DEFINE_string('config', None, 'Strategy config file.')

  logging.basicConfig(level='INFO', format='%(levelname)s %(asctime)s %(name)s] %(message)s')
  app.run(main)
